#!/usr/bin/env python3

"""
Minimal example showing how to use mcp.server.websocket.websocket_server
to implement a WebSocket-based MCP server endpoint using Starlette.

Notes:
- websocket_server is an async context manager that expects raw ASGI
  (scope, receive, send). Do NOT call websocket.accept() yourself:
  websocket_server will accept the connection with the "mcp" subprotocol.
- read_stream yields either SessionMessage or an Exception (ValidationError).
- write_stream expects SessionMessage objects whose .message is a pydantic model
  (the ws writer inside websocket_server calls .model_dump_json()).
- Run with an ASGI server (Hypercorn or Uvicorn). Hypercorn is recommended
  for anyio compatibility.
"""
import logging
import anyio

from starlette.applications import Starlette
from starlette.routing import WebSocketRoute
from starlette.websockets import WebSocket

from mcp.server.websocket import websocket_server
from mcp.shared.message import SessionMessage
import mcp.types as types  # MCP JSON-RPC types / models

logger = logging.getLogger("mcp-example")
logger.setLevel(logging.INFO)


async def mcp_endpoint(websocket: WebSocket):
    """
    WebSocket endpoint used by Starlette. We do NOT call websocket.accept() here;
    instead we hand the ASGI components to websocket_server which will accept.
    """
    # Pass the raw ASGI interfaces provided by the Starlette WebSocket object
    async with websocket_server(websocket.scope, websocket.receive, websocket.send) as (read_stream, write_stream):
        # run the handler. The websocket_server already runs internal reader/writer
        # tasks that move data between the socket and the anyio streams. Here we
        # run one task to process session messages coming from read_stream and
        # to send replies via write_stream.
        async with read_stream:
            async for item in read_stream:
                # item can be an Exception (ValidationError) if client sent invalid JSON-RPC
                if isinstance(item, Exception):
                    logger.exception("Received invalid message from client", exc_info=item)
                    # optionally send an error response here if you can associate it with an id
                    continue

                session_msg: SessionMessage = item
                client_message = session_msg.message
                logger.info("received client JSON-RPC message: %s", client_message)

                # Example routing: if the incoming message has an `id`, treat it as a request
                # and return a basic response object. Adjust to actual types in mcp.types.
                # The exact model names may vary; inspect mcp.types to use the right constructors.
                try:
                    # If message has an 'id' it's a request -> send a response
                    if getattr(client_message, "id", None) is not None:
                        # Build a JSON-RPC response pydantic model. The exact model class
                        # name/constructor depends on mcp.types; this is a general pattern.
                        response_obj = types.JSONRPCMessage.model_validate({
                            "jsonrpc": "2.0",
                            "id": client_message.id,
                            "result": {"echo_method": getattr(client_message, "method", None)}
                        })
                        await write_stream.send(SessionMessage(response_obj))
                    else:
                        # Notification: no reply required. You may want to process it.
                        logger.info("notification received (no reply): %s", getattr(client_message, "method", None))
                except Exception:
                    logger.exception("failed to construct/send response")


app = Starlette(routes=[WebSocketRoute("/mcp", mcp_endpoint, name="mcp")])

# To run:
#   uvicorn server:app --port 8000
#
# Javascript client example:
#   const ws = new WebSocket("ws://localhost:8000/mcp", "mcp");
#   ws.onmessage = (ev) => console.log("server:", ev.data);
#   ws.send(JSON.stringify({ jsonrpc: "2.0", id: 1, method: "hello" }));


